package _binary_tree

import common.TreeNode
import org.junit.Assert
import org.junit.Test
import common.printLevelOrder2

/**
 * ```
 * 814. 二叉树剪枝
 * 给你二叉树的根结点 root ，此外树的每个结点的值要么是 0 ，要么是 1 。
 * 返回移除了所有不包含 1 的子树的原二叉树。
 * 节点 node 的子树为 node 本身加上所有 node 的后代。
 *
 * 示例 1：
 * 输入：root = [1,null,0,0,1]
 * 输出：[1,null,0,null,1]
 * 解释：
 * 只有红色节点满足条件“所有不包含 1 的子树”。 右图为返回的答案。
 *
 * 示例 2：
 * 输入：root = [1,0,1,0,0,0,1]
 * 输出：[1,null,1,null,1]
 *
 * 示例 3：
 * 输入：root = [1,1,0,1,1,0,1,0]
 * 输出：[1,1,0,1,1,null,1]
 *
 * 提示：
 * 树中节点的数目在范围 [1, 200] 内
 * Node.val 为 0 或 1
 *
 * ```
 */
class leetcode_814 {
    @Test
    fun test_1() {
        val n0 = TreeNode(0)
        val n0_2 = TreeNode(0)
        val n1 = TreeNode(1)
        val n1_2 = TreeNode(1)
        val root = n1
        n1.right = n0
        n0.left = n0_2
        n0.right = n1_2
        val actual = pruneTree(root)
        val expected = arrayListOf(1, null, 0, null, 1)
        Assert.assertEquals(expected.toString(), printLevelOrder2(actual).toString())
    }


    @Test
    fun test_2() {
        val n0 = TreeNode(0)
        val n0_2 = TreeNode(0)
        val n0_3 = TreeNode(0)
        val n0_4 = TreeNode(0)
        val n1 = TreeNode(1)
        val n1_2 = TreeNode(1)
        val n1_3 = TreeNode(1)
        val root = n1

        n1.left = n0
        n1.right = n1_2

        n0.left = n0_2
        n0.right = n0_3

        n1_2.left = n0_4
        n1_2.right = n1_3

        val actual = pruneTree(root)
        val expected = arrayListOf(1, null, 1, null, 1)
        Assert.assertEquals(expected.toString(), printLevelOrder2(actual).toString())
    }

    @Test
    fun test_3() {
        val n0 = TreeNode(0)
        val n0_2 = TreeNode(0)
        val n0_3 = TreeNode(0)

        val n1 = TreeNode(1)
        val n1_2 = TreeNode(1)
        val n1_3 = TreeNode(1)
        val n1_4 = TreeNode(1)
        val n1_5 = TreeNode(1)

        val root = n1

        n1.left = n1_2
        n1.right = n0


        n1_2.left = n1_3
        n1_2.right = n1_4

        n0.left = n0_2
        n0.right = n1_5


        n1_5.left = n0_3


        val actual = pruneTree(root)
        val expected = arrayListOf(1,1,0,1,1,null,1)
        Assert.assertEquals(expected.toString(), printLevelOrder2(actual).toString())
    }

    private fun pruneTree(root: TreeNode?): TreeNode? {
        /**
        思路：二叉树的后序遍历
        https://leetcode.cn/problems/binary-tree-pruning/solutions/1686077/by-ac_oier-7me9
        递归处理左右子树，并将新的左右子树设置给root。
        当左右子树有一个非空或value 不等 0，则返回root，否则返回null
         */


        // 1 确认函数参数以及返回值
        // 2 确认终止条件
        if (null == root) {
            return null
        }

        // 3 确认单层遍历的逻辑: 左、右、中
        // 左
        root.left = pruneTree(root.left)

        // 右
        root.right = pruneTree(root.right)

        // 中
        return if (root.left != null || root.right != null || root.`val` != 0) {
            return root
        } else null
    }
}